package frege.data.RaList where

-- a basic implementation of Chris Okasaki's random access list

data RaList a = EOL | Node Int (Heap a) (RaList a)
data Heap a = Leaf a | Branch a (Heap a) (Heap a) 

empty :: RaList a
empty = EOL

null :: RaList a -> Bool
null EOL = true
null _ = false

singleton :: a -> RaList a
singleton x = Node 0 (Leaf x) EOL

cons :: a -> RaList a -> RaList a
cons x (Node sz1 hp1 (Node sz2 hp2 tl)) 
    | sz1 == sz2 = Node (2 * sz1 + 1) (Branch x hp1 hp2) tl
cons x xs = Node 1 (Leaf x) xs

head :: RaList a -> a
head (Node _ (Leaf x) _) = x
head (Node _ (Branch x _ _) _) = x
head _ = error "EOL"

tail :: RaList a -> RaList a
tail (Node 1 _ tl) = tl
tail (Node sz (Branch _ left right) tl) = 
   Node sub left (Node sub right tl) where sub = sz `div` 2
tail _ = error "EOL"

size :: RaList a -> Int
size (Node sz _ tl) = sz + size tl 
size _ = 0

get :: RaList a -> Int -> a
get (Node sz hp tl) i | i < sz = getHeap sz hp i
                      | otherwise = get tl (i-sz)
get _ _ = error "EOL"

getHeap :: Int -> Heap a -> Int -> a 
getHeap _ (Leaf x) 0 = x
getHeap _ (Branch x _ _) 0 = x
getHeap sz (Branch _ left right) k = 
  if k <= sub then getHeap sub left (k-1) 
              else getHeap sub right (k-sub-1) where                       
  sub = sz `div` 2

toList :: RaList a -> [a]
toList (Node _ hp tl) = addHeapToList hp $ toList tl 
toList EOL = []

addHeapToList :: Heap a -> [a] -> [a]
addHeapToList (Leaf x) xs = x:xs
addHeapToList (Branch x left right) xs = x : addHeapToList left (addHeapToList right xs)

fromList :: [a] -> RaList a
fromList xs = Prelude.foldr cons EOL xs

map :: (a -> b) -> RaList a -> RaList b
map f (Node sz hp tl) = Node sz (mapHeap f hp) (map f tl) where
   mapHeap f (Leaf x) = Leaf (f x)
   mapHeap f (Branch x left right) = Branch (f x) (mapHeap f left) (mapHeap f right)  
map _ _ = EOL

drop :: Int -> RaList a -> RaList a
drop n EOL = EOL
drop n (node @ Node sz hp tl) 
    | n <= 0 = node
    | n >= sz = drop (n - sz) tl
drop n (Node sz (Branch _ left right) tl) = drop (n - 1) (Node sub left (Node sub right tl)) where
   sub = sz `div` 2 
   
take :: Int -> RaList a -> RaList a
take _ EOL = EOL  
take n _ | n <= 0 = EOL
take n (Node sz hp tl) 
  | sz >= n = fromList $ Prelude.take n $ addHeapToList hp []
  | otherwise = addHeap hp $ take (n-sz) tl where 
     addHeap (Leaf x) xs = cons x xs
     addHeap (Branch x left right) xs = cons x (addHeap left (addHeap right xs)) 

foldl :: (a -> b -> a) -> a -> RaList b -> a 
foldl f x EOL = x
foldl f x xs = foldl f (f x $ head xs) (tail xs)

foldr :: (a -> b -> b) -> b -> RaList a -> b 
foldr f x EOL = x
foldr f x (Node _ hp tl) = foldrHeap f (foldr f x tl) hp where
   foldrHeap f x (Leaf y) = f y x
   foldrHeap f x (Branch y left right) = f y $ foldrHeap f (foldrHeap f x right) left

filter :: (a -> Bool) -> RaList a -> RaList a
filter cond EOL = EOL
filter cond xs = let hd = head xs
                     tl = filter cond $ tail xs
                 in if cond hd then cons hd tl else tl   

zipWith :: (a -> b -> c) -> RaList a -> RaList b -> RaList c
zipWith f EOL _ = EOL
zipWith f _ EOL = EOL
zipWith f xs ys = cons (f (head xs) (head ys)) $ zipWith f (tail xs) (tail ys)

zip :: RaList a -> RaList b -> RaList (a,b)
zip xs ys = zipWith (,) xs ys

unzip :: RaList (a,b) -> (RaList a, RaList b)
unzip xs = (map fst xs, map snd xs)

instance Show Show a => RaList a where
  show ra = show $ RaList.toList ra
